#ifndef TREE_H
#define TREE_H

#include <math.h>
#include <vector>
#include <string>
#include <iostream>
#include <sstream>
#include <fstream>
#include <stdlib.h>
#include <stdio.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_rng.h>

const double epsilon = 1e-12;
const double pi = 3.14159265358979323846;
const double e = 2.71828182845904523536;
const double p_1 = 3;
const double p_2 = 1;
gsl_rng * r;

using namespace::std;

// struct node forms the nodes and branches within the tree.  

struct node
{

	int right;
	int left;
	int parent;

	double rightBranch;
	double leftBranch;
	double parentBranch;
//	double cone;
	
	node() 
	{ right = -1; left = -1; parent = -1; rightBranch = -1; parentBranch = -1; leftBranch = -1; }//cone = -1;}

	void operator=(node rhs) 
	{ 
	  right = rhs.right;  left = rhs.left; parent = rhs.parent; 
	  rightBranch = rhs.rightBranch; leftBranch = rhs.leftBranch; parentBranch = rhs.parentBranch; 
//	  cone = rhs.cone;
	}
};


// assignment overloaded for node

ostream& operator << (ostream & os, const node& s)
{  
	    os<<"Entries: Left: "<<s.left<<"\t Right: "<<s.right<<"\t Parent: "<<s.parent<<endl
		<<"Lengths: Left: "<<s.leftBranch<<"\t Right: "<<s.rightBranch<<"\t Parent: "<<s.parentBranch<<endl;
//		<<"Cone: "<<s.cone<<endl;
}



//class tree
	// basic tree class 
	// includes input/output, 
	// generates randomTree, normalizes tree, displays tree, tell whether an entry is a tip

class tree
{
	public:
		//basic variables
		vector<node> treeNodes;
		vector<int> lineageNodeHash;
		vector<int> nodeLineageHash;

		//constructors...etc.
		tree(){ treeNodes.clear(); lineageNodeHash.clear(); nodeLineageHash.clear(); r = gsl_rng_alloc(gsl_rng_rand48); 
			long seed; seed = rand()*time (NULL) * getpid();    gsl_rng_set (r, seed);  };	
		~tree(){};
	
		tree & operator = (const tree & rhs)
		{
			if (this != &rhs)
			{
			    treeNodes = rhs.treeNodes; lineageNodeHash = rhs.lineageNodeHash; nodeLineageHash = rhs.nodeLineageHash;
			}
			return *this;
	    	}


		// Displays the tree, node by node
		void displayTree(void);		   // DONE
		// Generates an unnormalized, random tree.
		int generateRandomTree(int k);	 // DONE
		// Normalizes the tree branch lengths by Z
		void normalizeSubtree(double Z, int);   // DONE
		double findPathAbove(int node);
		void makeTreeUltrameric();
		void normalizeTree();
		int randomNodeSPR();
		int randomNodeNotSPR();
		vector <int> findNodesBelow(int);
		double findPathBelow(int);
 		// Finds the tree length below eswapNode(int node1, int node2)ntry
		double subTreeLength(int nodeEntry);  // DONE
		// Gives boolean: true if entry is a tip, false otherwise
		int isTip(int);			 // DONE
		// Returns all leaf partitions of the tree
		vector< vector<double> > getAllPartitions();  
		// Sets lineageNodeHash so entries either -1 or corresponding lineage
		void hashLN(void);  
		// Returns vector of all tips below a particular node;
		vector<int> findTipsBelow(int entry); // DONE
		// Returns a random node
		int randomNode();  			//DONE
		// Changes the parent branch length to a given node, to a given length
		void changeParentBranchLength(int, double);  //DONE
		// Returns 1 = right, 2 = left , 3 = parent
		int getParentNodeDirection(int);	//DONE
		// Returns the total length of the tree
		double totalLengthTree(void);
		//	
		void recurseUltra(int);
		void detachNode(int nodeRemove, int otherNode);
		// Implements a random SPR move
		int randomSPR(void);
		// Removes a pair of nodes
		// Returns the number corresponding to the lineageNodeHash
		int nodeNumber(int);
		string outputTree();
		void attachNode(int,int);
		int treeConsistency();
		double nodeConsistency(int);

	private:
		// Changes the parent of a node to a new specified parent
		void replaceValue(int,int,int);
		// Places a new node
		void placeNode(int, double, double, int, int);			
		// Insert a node a given split
		void insertNode(int, int);
		// Returns a random node not below 
		int randomNodeNotBelow(int node1, int node2);
		int randomOtherNode(int node);
		int insertRandomNode(int);
		void swapNode(int, int);
		int checkRoot(void);
		int isRootBranch(int ,int );
		int isInSet(vector<int>,int);
		void swapRoot(int);
		void reformNode(int);
		void stripNode(int);
		void replaceBranch(int, int, double);
		int whatLastNumber(vector<int> );
		string recurseTree(int node);

};

ostream & operator << (ostream & os, const tree& s)
{
	int N = s.treeNodes.size(), i;

	for (i =0; i < N; i++)
	{	
		os<<i<<endl;
		os<<s.treeNodes[i];	}
	os<<endl;
}

//ostream& operator << (ostream & os, const node& s)

double tree::nodeConsistency(int node)
{
	double return_val = 0, total = 0;
	if (node !=0)
	{
		if (treeNodes[treeNodes[node].parent].right == node)
		{
			total++;
			return_val += (double)(treeNodes[treeNodes[node].parent].rightBranch == treeNodes[node].parentBranch);
		}
		else
		{
			total++;
			return_val += (double)(treeNodes[treeNodes[node].parent].leftBranch == treeNodes[node].parentBranch);
		}
		if (treeNodes[node].parent!=-1)
		{
			total++;
			return_val += (double)(treeNodes[node].parentBranch>0);
		}

		if (treeNodes[node].right!=-1)
		{
			total++;
			return_val += (double)(treeNodes[node].rightBranch>0);
		}

		if (treeNodes[node].left!=-1)
		{
			total++;
			return_val += (double)(treeNodes[node].leftBranch>0);
		}
	//	if (isTip(node))
	//	{
	//		total++;
	//		return_val += (double)(treeNodes[node].cone > 0);
	//	}
	}
	else
	{
		return(1.0);
	}
	return(return_val/total);
}

int tree::treeConsistency()
{
	int i =0;
	double totalPath=0, currentPath, numTips = ((double)treeNodes.size()+1)/2;
	for (i = 0; i< treeNodes.size();i++)
	{
		if (isTip(i))
		{
			currentPath = findPathAbove(i);
			totalPath += currentPath;
		}
	}
	if (fabs(currentPath*numTips - totalPath)>1e-6)
	{
		cout<<currentPath<<" "<<totalPath<<" "<<numTips<<" "<<fabs(currentPath*numTips - totalPath)<<endl;
		cout<<"Inconsistent path lengths...";
		exit(-1);
	}
	for (i = 0; i<treeNodes.size();i++)
	{
	//	cout<<"Node "<<i<<" consistent? "<<nodeConsistency(i)<<endl; 
		if (nodeConsistency(i)<1)
		{
			cout<<"Node inconsistent...";
			exit(-1);	
		}
	}
	return(1);
}

string tree::recurseTree(int node)
{
	ostringstream oss;

	if (isTip(node)==0)
	{
		if (node == 0)
		{
			oss << "(" << recurseTree(treeNodes[node].right) <<","<<recurseTree(treeNodes[node].left)<<")";
		}
		else
		{
			oss << "(" << recurseTree(treeNodes[node].right) <<","<<recurseTree(treeNodes[node].left)<<"):"<<treeNodes[node].parentBranch;
		}
	}
	else
	{
		oss<<node<<":"<<treeNodes[node].parentBranch;
	}

	return oss.str();
}

string tree::outputTree()
{
	
	int currentPos, check = 1, numBranch = treeNodes.size()-1, counter =0,counter_2, counter_3,check_2, num_zeros;

	ostringstream oss;

	oss <<recurseTree(0) <<";";

//	cout<<oss.str()<<endl;	
	return oss.str();
}





double tree::totalLengthTree(void)
{
	double total = 0;
	int i = 0;

	total = subTreeLength(0);
	return total;	
}

int tree::randomNodeSPR()
{
	int sample;

	int check=1, right = treeNodes[0].right, left = treeNodes[0].left;
	
	while(check)
	{
		sample = floor(gsl_ran_flat(r,0,1)*treeNodes.size());
		if ( (sample != (int)treeNodes.size()) && (sample!=0) &&(treeNodes[0].right != sample) && (treeNodes[0].left != sample) )
		{	check = 0;	}			
	}

	return sample;
}


int tree::randomNodeNotSPR()
{
	int sample;

	int check=1, right = treeNodes[0].right, left = treeNodes[0].left;
	
	while(check)
	{
		sample = floor(gsl_ran_flat(r,0,1)*treeNodes.size());
		if ( (sample != (int)treeNodes.size()) && (sample!=0) )
		{	check = 0;	}			
	}

	return sample;
}

void tree::hashLN(void)
{
	int i = 0, counter = 0;

	lineageNodeHash.clear();
	lineageNodeHash.resize(treeNodes.size(),-1);

	nodeLineageHash.clear();
	nodeLineageHash.resize((treeNodes.size()+1)/2,-1);
	for (i =0; i<(int)treeNodes.size();i++)
	{
		if(isTip(i))
		{
			nodeLineageHash[counter] = i;
			lineageNodeHash[i] = counter; 	counter++;
			
		}
	}


}

int tree::isTip(int test)
{
	if (test < (int)treeNodes.size())
	{
		if ((treeNodes[test].right == -1)&&(treeNodes[test].left == -1))
		{
			return 1;
		}
		else
		{	return 0; 	}
	}
	else
	{
		cerr<<"Node does not exist!"<<endl;
		exit(-1);
	}
}

void tree::displayTree()
{
	int i = 0;

	cout<<"The tree contains "<<treeNodes.size()<<" nodes."<<endl;
	for (i = 0; i<(int)treeNodes.size();i++)
	{
		cout<<"Node "<<i<<endl;
		cout<<treeNodes[i]<<endl;
	}
}

double tree::findPathAbove(int node)
{
	if (node == 0)
	{	return(0);	}
	else
	{	
		if (isTip(node))
		{	return(treeNodes[node].parentBranch + findPathAbove(treeNodes[node].parent));}//treeNodes[node].cone;	}
		else
		{	return(treeNodes[node].parentBranch + findPathAbove(treeNodes[node].parent));	}	
	}
}

void tree::makeTreeUltrameric()
{
	int i;
	double treeLength = subTreeLength(0), abovePath;

//	treeLength = subTreeLength(0);
//	normalizeSubtree(treeLength,0);



	//cout<<treeLength<<endl<<endl;
	recurseUltra(0);
		
	//treeLength = subTreeLength(0);

	//normalizeSubtree(treeLength,0);
}

void tree::recurseUltra(int node)
{
	int parent;
	double abovePath = 0,test, totalBranch, split;
	
	if (node == 0)
	{
		recurseUltra(treeNodes[node].right);
		recurseUltra(treeNodes[node].left);
	}
	else
	{
		parent = treeNodes[node].parent;
		abovePath = findPathAbove(parent);

//		cout<<node<<" "<<parent<<" "<<abovePath<<endl;
		if (isTip(node))
		{	

			totalBranch = 1-abovePath;
			treeNodes[node].parentBranch = totalBranch;
//			treeNodes[node].cone = (1-split)*totalBranch;

			if (treeNodes[parent].right == node)
			{	treeNodes[parent].rightBranch  = treeNodes[node].parentBranch;	}
			else
			{	treeNodes[parent].leftBranch  = treeNodes[node].parentBranch;	}
		}
		else
		{
			treeNodes[node].parentBranch = (1-abovePath)/2;
			if (treeNodes[node].parentBranch < 0)
			{
				cout<<treeNodes[node].parentBranch<<endl;
				exit(-1);
			}
			if (treeNodes[parent].right == node)
			{	treeNodes[parent].rightBranch  = treeNodes[node].parentBranch;	}
			else
			{	treeNodes[parent].leftBranch  = treeNodes[node].parentBranch;	}

			recurseUltra(treeNodes[node].right);
			recurseUltra(treeNodes[node].left);
		}
	}
}

int tree::generateRandomTree(int k)
{
	int size = k;

	treeNodes.clear();

	node dummy;
	treeNodes.push_back(dummy);
	double length = 0;

	while(size>1)
	{
		insertRandomNode(0);
		size--;
	}

//	displayTree();
	makeTreeUltrameric();
//	displayTree();
	return 1;

}

int tree::insertRandomNode(int entry)
{
	int total = (treeNodes[entry].right < 0) + (treeNodes[entry].left<0)+(treeNodes[entry].parent<0);
	double test;
	node dummy1, dummy2;

	srand(time(NULL)-rand());
	double dummy_rand;
	
	
	if ((total == 3)) // then there is nothing there
	{   
		treeNodes[entry].left = treeNodes.size();
		treeNodes[entry].right = treeNodes.size()+1;

		dummy1.parent = entry;
		dummy2.parent = entry;
	 	dummy1.parentBranch = 0.5;//gsl_ran_flat(r,1,5); 
		dummy2.parentBranch = 0.5;//gsl_ran_flat(r,1,5); 

		treeNodes[entry].leftBranch = dummy1.parentBranch; treeNodes[entry].rightBranch = dummy2.parentBranch; 
		treeNodes.push_back(dummy1); treeNodes.push_back(dummy2); 
		return 1;		
	}
	if (total == 1) // then this is the root and the parent is open
	{
		// pick a random branch and go down
		// if root, three options; if not then, two
		dummy_rand = (entry % 2);//gsl_ran_flat(r,0,2);
		if (dummy_rand == 1)
		{	return insertRandomNode(treeNodes[entry].right);  }
		else
		{	return insertRandomNode(treeNodes[entry].left);  }
			
		return 1;
	}
	if (total == 0) // then an internal node;
	{
		// pick a random branch and go down
		// if root, three options; if not then, two
		dummy_rand = entry%3;//gsl_ran_flat(r,0,3);

		if (entry == 0)
		{
			if (dummy_rand  >= 2)
			{
				return insertRandomNode(treeNodes[entry].parent);
			}
			else
			{
				if (dummy_rand <= 1)
				{	return insertRandomNode(treeNodes[entry].left);    }
				else
				{	return insertRandomNode(treeNodes[entry].right);   }
			}
		}
		else
		{
			dummy_rand = entry%3;//gsl_ran_flat(r,0,2);
			if (dummy_rand >= 1 )
			{	return insertRandomNode(treeNodes[entry].right);  }
			else
			{	return insertRandomNode(treeNodes[entry].left);  }
		}
		return 1;			
	}
	if (total == 2)
	{
		treeNodes[entry].left = treeNodes.size();
		treeNodes[entry].right = treeNodes.size()+1;
//		treeNodes[entry].cone = -1;

		dummy1.parent = entry;
		dummy2.parent = entry;
	 	dummy1.parentBranch = 0.5;//gsl_ran_flat(r,1,5); 
		dummy2.parentBranch = 0.5;//gsl_ran_flat(r,1,5); 
//		dummy1.cone = gsl_ran_flat(r,1,5);
//		dummy2.cone = gsl_ran_flat(r,1,5);


		treeNodes[entry].leftBranch = dummy1.parentBranch; treeNodes[entry].rightBranch = dummy2.parentBranch; 
		treeNodes.push_back(dummy1); treeNodes.push_back(dummy2); 
		return 1;	
	}
		
	
	
}


double tree::findPathBelow(int entry)
{
	if (isTip(entry))
	{
		return 0;//treeNodes[entry].cone;
	}
	else
	{
		if (treeNodes[entry].left!=-1)
		{
			return findPathBelow(treeNodes[entry].left) + treeNodes[entry].leftBranch;
		}
		else
		{
			return findPathBelow(treeNodes[entry].right) + treeNodes[entry].rightBranch;
		}
	}
}

//Divides all branch lengths on subtree by the quantity length

void tree::normalizeSubtree(double length, int nodeEntry)
{
	int i = 0;
	
	if (treeNodes[nodeEntry].right !=-1)
	{	
		treeNodes[nodeEntry].rightBranch = treeNodes[nodeEntry].rightBranch/length;	
		treeNodes[treeNodes[nodeEntry].right].parentBranch = treeNodes[treeNodes[nodeEntry].right].parentBranch/length;
		normalizeSubtree(length,treeNodes[nodeEntry].right);
	}
	if (treeNodes[nodeEntry].left !=-1)
	{	
		treeNodes[nodeEntry].leftBranch = treeNodes[nodeEntry].leftBranch/length;
		treeNodes[treeNodes[nodeEntry].left].parentBranch = treeNodes[treeNodes[nodeEntry].left].parentBranch/length;
		normalizeSubtree(length,treeNodes[nodeEntry].left);
	}
//	if (treeNodes[nodeEntry].cone !=-1)
//	{
//		treeNodes[nodeEntry].cone = treeNodes[nodeEntry].cone/length;
//	}
	
	
}

// Takes node entry and finds the length of the tree below it 
// Returns the total;
// Works recursively
double tree::subTreeLength(int nodeEntry)
{
	double total = 0;
	if (treeNodes[nodeEntry].right != -1)
	{	total = total + treeNodes[nodeEntry].rightBranch + subTreeLength(treeNodes[nodeEntry].right);	}
	if (treeNodes[nodeEntry].left != -1)
	{	total = total + treeNodes[nodeEntry].leftBranch + subTreeLength(treeNodes[nodeEntry].left);	}	
//	if (isTip(nodeEntry))
//	{	total = total + treeNodes[nodeEntry].cone;	}
	return total;
}


vector<int> tree::findTipsBelow(int entry)
{
	vector<int> nodes;
	vector<int> right;
	vector<int> left;
	nodes.clear();
	right.clear();
	left.clear();
	int i = 0;

	if (!isTip(entry))
	{
		
		if (treeNodes[entry].left != -1)
		{
			left = findTipsBelow(treeNodes[entry].left); 
			for (i = 0; i <(int)left.size();i++)
			{   nodes.push_back(left[i]); };
		}
		if (treeNodes[entry].right != -1)
		{	
			right = findTipsBelow(treeNodes[entry].right); 
			for (i = 0; i <(int)right.size();i++)
			{   nodes.push_back(right[i]); };
		}
	}
	else
	{
		nodes.push_back(entry);
	}
	return nodes;
}

vector<int> tree::findNodesBelow(int entry)
{
	vector<int> nodes;
	vector<int> right;
	vector<int> left;
	vector<int> parent;
	nodes.clear();
	right.clear();
	left.clear();
	parent.clear();
	int i = 0;

	nodes.push_back(entry);
	if (!isTip(entry))
	{	
		
		if (treeNodes[entry].right !=-1)
		{	
			right = findNodesBelow(treeNodes[entry].right);
			for (i = 0; i<right.size();i++)
			{
				nodes.push_back(right[i]);
			}		
		}
		if (treeNodes[entry].left !=-1)
		{
			left = findNodesBelow(treeNodes[entry].left);

			for (i = 0; i<left.size();i++)
			{
				nodes.push_back(left[i]);
			}	
		}		
		
	}
	return nodes;
}



// Takes nothing
// Finds all of the possible partitions for a given tree
// 

int tree::randomSPR(void)
{

	// grab a random node that is not the root or its children

	int randomNode = randomNodeSPR(), parentRandom, check = 1, attachNode, c1, c2, counter,i;

	parentRandom = treeNodes[randomNode].parent;
	c1 = treeNodes[treeNodes[randomNode].parent].right;
	c2 = treeNodes[treeNodes[randomNode].parent].left;	
	// take the parent branch and detach it
	
	detachNode(randomNode,parentRandom);

	// remember the branch that is altered so no further contact	
	vector <int> set = findNodesBelow(parentRandom);

	//displayTree();

	// grab a random node, not at the base of the altered branch
	counter = 0;
	while(check)
	{
		counter++;
		if (counter > 100)
		{	exit(-1);	}
		attachNode = randomNodeNotSPR();
		if (isInSet(set, attachNode)==0)
		{
			check = 0;
		}
	}

	//displayTree();
	//cout<<attachNode<<endl;
	// attach on the parent branch
	insertNode(parentRandom,attachNode);

}


void tree::swapNode(int node1, int node2)
{
	node dummy;
	
	dummy = treeNodes[node1];
	treeNodes[node1] = treeNodes[node2];
	treeNodes[node2] = dummy;
}

void tree::stripNode(int node1)
{
	treeNodes[node1].parent = -1;
	treeNodes[node1].left = -1;
	treeNodes[node1].right = -1;
	treeNodes[node1].parentBranch = -1;
	treeNodes[node1].leftBranch = -1;
	treeNodes[node1].rightBranch = -1;
//	treeNodes[node1].cone = -1;
}

void tree::insertNode(int nodeToAttach, int nodeToBe)
{
	
	int parentToBe = treeNodes[nodeToBe].parent;

	double branchBelow, branchAbove, split = gsl_ran_flat(r,epsilon,1-epsilon), totalLength, lengthAbove, lengthBelow, scale;

	branchBelow = treeNodes[nodeToBe].parentBranch*(1-split);
	branchAbove = treeNodes[nodeToBe].parentBranch*(split);

	totalLength = findPathBelow(0);
	lengthAbove = findPathAbove(parentToBe) + branchAbove;
	lengthBelow = findPathBelow(nodeToAttach);
	scale = lengthBelow/(totalLength-lengthAbove);
	normalizeSubtree(scale,nodeToAttach);

	treeNodes[nodeToAttach].parent = parentToBe;
	treeNodes[nodeToAttach].parentBranch = branchAbove;

	if (treeNodes[parentToBe].right == nodeToBe)
	{
		treeNodes[parentToBe].right = nodeToAttach;	
		treeNodes[parentToBe].rightBranch = branchAbove;
	}
	else
	{
		treeNodes[parentToBe].left = nodeToAttach;	
		treeNodes[parentToBe].leftBranch = branchAbove;
	}

	if (treeNodes[nodeToAttach].right == -1)
	{
		treeNodes[nodeToAttach].right = nodeToBe;
		treeNodes[nodeToAttach].rightBranch = branchBelow;
		treeNodes[nodeToBe].parent = nodeToAttach;
		treeNodes[nodeToBe].parentBranch = branchBelow;
	}
	else
	{
		treeNodes[nodeToAttach].left = nodeToBe;
		treeNodes[nodeToAttach].leftBranch = branchBelow;
		treeNodes[nodeToBe].parent = nodeToAttach;
		treeNodes[nodeToBe].parentBranch = branchBelow;
	}


}



int tree::isInSet(vector<int> set, int val)
{
	int check = 0;
	for (int i = 0; i < set.size(); i++)
	{
		if (set[i] ==val)
		{	check = 1;	}
	}
	return check;
}

void tree::detachNode(int nodeRemove, int parentRemove)
{
	// assumes that parentRemove is the parents of nodeRemove

	int otherBelow;

//	cout<<"TOGO:"<<nodeRemove<<" "<<parentRemove<<endl;
	double branch = 0;

	if (treeNodes[parentRemove].right == nodeRemove)
	{
		otherBelow = treeNodes[parentRemove].left;	
		treeNodes[parentRemove].left = -1;		
		treeNodes[parentRemove].leftBranch = -1;
	}
	else
	{	
		otherBelow = treeNodes[parentRemove].right;	
		treeNodes[parentRemove].right = -1;
		treeNodes[parentRemove].rightBranch = -1;
	}

	if (treeNodes[treeNodes[parentRemove].parent].right == parentRemove)
	{
		branch = treeNodes[otherBelow].parentBranch + treeNodes[parentRemove].parentBranch;
		treeNodes[treeNodes[parentRemove].parent].right = otherBelow;
		treeNodes[treeNodes[parentRemove].parent].rightBranch = branch;
		treeNodes[otherBelow].parent = treeNodes[parentRemove].parent;
		treeNodes[otherBelow].parentBranch = branch;
	}
	else
	{
		branch = treeNodes[otherBelow].parentBranch + treeNodes[parentRemove].parentBranch;
		treeNodes[treeNodes[parentRemove].parent].left = otherBelow;
		treeNodes[treeNodes[parentRemove].parent].leftBranch = branch;
		treeNodes[otherBelow].parent = treeNodes[parentRemove].parent;
		treeNodes[otherBelow].parentBranch = branch;
	}

	treeNodes[parentRemove].parent = -1;
	treeNodes[parentRemove].parentBranch = -1;
}




#endif /* TREE_H */
